import pandas as pd
from gm.api import get_instruments, get_history_instruments, current, history_n
from gm.model.storage import context

from ...common.date import jue_jin_get_nearly_trading_date, is_same_day


# from gmcache import *


# @jit(nopython=True)
def get_stocks_can_trade_today(context, skip_st):
    """
    获取最新的今天可交易的上市公司股票代码，自动忽略688开头的科创板股票和停牌股，可选择性忽略ST股
    :param context:
    :param skip_st:是否忽略st股
    :return: 控制返回空
    """
    nearly_trading_date = jue_jin_get_nearly_trading_date(context)
    # 获取最新的交易标的信息数据,回测返回的是系统当前天对应的交易标的数据，回测不可用
    all_stock = get_instruments(exchanges='SHSE, SZSE', sec_types=[1],
                                skip_suspended=False,
                                skip_st=False, df=True)
    if len(all_stock) == 0:
        return all_stock

    # 实盘模式
    if context.mode == 1:
        # 剔除B股、科创板股票和停牌股
        all_stock = all_stock[(all_stock['symbol'].str[5] != '9') & (
                all_stock['symbol'].str[5] != '2') & (
                                  ~ all_stock['symbol'].str.startswith('SHSE.688')) & (
                                      all_stock['trade_date'] == nearly_trading_date) & (
                                      all_stock['delisted_date'] >= nearly_trading_date)]
    # 回测模式
    if context.mode == 2:
        # 获取最新的交易标的信息数据,回测返回的是系统当前天对应的交易标的数据，非回测不可用
        all_stock = get_history_instruments(symbols=all_stock['symbol'].tolist(), start_date=nearly_trading_date,
                                            end_date=nearly_trading_date,
                                            df=True)
        # 剔除B股和科创板股票
        all_stock = all_stock[(all_stock['symbol'].str[5] != '9') & (
                all_stock['symbol'].str[5] != '2') & (
                                  ~ all_stock['symbol'].str.startswith('SHSE.688')) & (
                                      all_stock['trade_date'] == nearly_trading_date) & (
                                      all_stock['delisted_date'] >= nearly_trading_date)]

    # 剔除st股
    if skip_st and len(all_stock) > 0:
        all_stock = all_stock[all_stock['sec_level'] == 1]

    return all_stock


def get_his_n_data(symbol, frequency, count, end_time, adjust, with_new_tick_price):
    history_n_data = history_n(symbol=symbol, frequency=frequency, count=count, end_time=end_time,
                               fields='symbol, open, close, high, low, eob', adjust=adjust, df=True)
    if with_new_tick_price:
        # 实盘模式
        if context.mode == 1:
            current_data = current(symbols=symbol)
            last_tick_date = current_data[0].created_at
            last_his_data = history_n_data.loc[len(history_n_data) - 1]
            last_hist_date = last_his_data.eob
            if not is_same_day(last_tick_date, last_hist_date):
                df = pd.DataFrame(
                    [current_data[0]['symbol'], current_data[0]['open'], current_data[0]['price'],
                     current_data[0]['high'],
                     current_data[0]['low'], current_data[0]['created_at']]
                ).T
                df.columns = history_n_data.columns
                history_n_data = pd.concat([history_n_data, df], axis=0, ignore_index=True)
            if is_same_day(last_tick_date, last_hist_date):
                history_n_data.loc[len(history_n_data) - 1] = [current_data[0]['symbol'], current_data[0]['open'],
                                                               current_data[0]['price'], current_data[0]['high'],
                                                               current_data[0]['low'], current_data[0]['created_at']]

    return history_n_data
